{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "f40becec",
   "metadata": {},
   "outputs": [],
   "source": [
    "import onnxruntime as ort\n",
    "model_path = \"CSOL-Utilities-resnet18-800x600.onnx\"\n",
    "session = ort.InferenceSession(model_path)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "62285e3b",
   "metadata": {},
   "outputs": [],
   "source": [
    "from torchvision import transforms\n",
    "width = 800\n",
    "height = 600\n",
    "transform = transforms.Compose([\n",
    "    transforms.Pad(padding=20, fill=0, padding_mode='constant'),\n",
    "    transforms.Resize((height, width)),\n",
    "    transforms.ToTensor(),\n",
    "    transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),\n",
    "])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "111ac5f2",
   "metadata": {},
   "outputs": [],
   "source": [
    "from PIL import Image\n",
    "import numpy as np\n",
    "import torch\n",
    "\n",
    "def predict(image_path):\n",
    "    image = Image.open(image_path).convert(\"RGB\")\n",
    "    input_tensor = transform(image).unsqueeze(0)\n",
    "    input_numpy = input_tensor.numpy()\n",
    "    input_name = session.get_inputs()[0].name\n",
    "    output_name = session.get_outputs()[0].name\n",
    "    outputs = session.run([output_name], {input_name: input_numpy})\n",
    "    output_tensor = torch.tensor(outputs[0])\n",
    "    labels = [\n",
    "        \"大厅界面\",\n",
    "        \"等候室界面\",\n",
    "        \"加载界面\",\n",
    "        \"场景界面\",\n",
    "        \"结算界面\"\n",
    "    ]\n",
    "    # softmax\n",
    "    output_tensor = torch.nn.functional.softmax(output_tensor, dim=1)\n",
    "    for i, label in enumerate(labels):\n",
    "        print(f\"{label}: {output_tensor[0, i].item():.4f}\")\n",
    "    predicted_index = torch.argmax(output_tensor, dim=1).item()\n",
    "    predicted_value = output_tensor[0, predicted_index].item()\n",
    "    if (predicted_value < 0.95):\n",
    "        print(\"Result: 未知\")\n",
    "    else:\n",
    "        print(f\"Result: {labels[int(predicted_index)]}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "5d9ef733",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "大厅界面: 1.0000\n",
      "等候室界面: 0.0000\n",
      "加载界面: 0.0000\n",
      "场景界面: 0.0000\n",
      "结算界面: 0.0000\n",
      "Result: 大厅界面\n"
     ]
    }
   ],
   "source": [
    "p = \"C:/Users/Silver/BaiduSyncdisk/Pictures/Screen Shots/2025/2025-10-07_11-46-05_527.png\"\n",
    "predict(p)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.10.18"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
